# -*- coding: UTF-8 -*-
from gensim.test.utils import datapath
from gensim.models import KeyedVectors
from gensim import utils
import gensim.models
import tempfile

class MyCorpus:
    """An iterator that yields sentences (lists of str)."""

    def __iter__(self):
        # sanguo_nlp.txt 为分词后的三国演义文本
        corpus_path = datapath('C:/code/word2vce/py/sanguo_nlp.txt')
        for line in open(corpus_path):
            # assume there's one document per line, tokens separated by whitespace
            yield utils.simple_preprocess(line)

sentences = MyCorpus()
model = gensim.models.Word2Vec(sentences=sentences)

word_vectors = model.wv
# word_vectors : key: \u ,value: float

# 存储模型
word_vectors.save('C:/code/word2vce/py/sanguo02.kv')

# print(word_vectors)
# 查看向量
# print(word_vectors.vectors)
# 查看词和对应向量
# print(word_vectors.vocab)

# 查找相似词
result = word_vectors.most_similar(positive=[unicode('刘备','utf-8'), unicode('刘皇叔','utf-8'),unicode('玄德','utf-8')], negative=[unicode('曹操','utf-8')])

print("positive is :" + "刘备 刘皇叔 玄德 ." + " negative is : " +  "曹操")
print("result:")
for i in range(10) :
    most_similar_key, similarity = result[i]
    # 打印相似值
    print("Top " + str(i) + ":")
    print(unicode.encode(most_similar_key,'utf-8') + "   " + str(similarity))


result2 = word_vectors.wv[unicode('啼哭','utf-8')]
print("-------------------")
print(result2[0])

# 加载模型
# reloaded_word_vectors = KeyedVectors.load('C:/code/word2vce/py/sanguo02.kv')
